package main

import (
	"bufio"
	"encoding/base64"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"
)

// Proxy 定义了一个 HTTP/HTTPS 代理服务器的结构体。
type Proxy struct {
	localAddr   string            // 代理服务器的本地监听地址
	listener    net.Listener      // TCP 监听器
	wg          sync.WaitGroup    // WaitGroup 用于等待所有 goroutine 完成
	AuthEnabled bool              // 控制是否开启代理身份验证的标志
	Credentials map[string]string // 存储合法的用户名和密码 (username:password)
}

// NewProxy 创建并返回一个新的 Proxy 实例。
func NewProxy(localAddr string, authEnabled bool, credentials map[string]string) (*Proxy, error) {
	return &Proxy{
		localAddr:   localAddr,
		AuthEnabled: authEnabled,
		Credentials: credentials,
	}, nil
}

// Start 启动代理服务。
func (p *Proxy) Start() error {
	listener, err := net.Listen("tcp", p.localAddr)
	if err != nil {
		return err
	}
	p.listener = listener
	log.Printf("HTTP/HTTPS 代理正在监听 %s", p.localAddr)
	if p.AuthEnabled {
		log.Println("代理身份验证已启用")
	} else {
		log.Println("代理身份验证已禁用")
	}

	for {
		conn, err := p.listener.Accept()
		if err != nil {
			log.Printf("接受连接失败: %v", err)
			continue
		}
		p.wg.Add(1)
		go p.handleConnection(conn)
	}
}

// checkAuth 验证客户端提供的凭据。
func (p *Proxy) checkAuth(req *http.Request) (string, bool) {
	authHeader := req.Header.Get("Proxy-Authorization")
	if authHeader == "" {
		return "", false // 没有提供凭据
	}

	// 凭据通常以 "Basic " 开头
	if !strings.HasPrefix(authHeader, "Basic ") {
		return "", false
	}

	// 解码 Base64 编码的凭据
	encodedCredentials := strings.TrimPrefix(authHeader, "Basic ")
	decoded, err := base64.StdEncoding.DecodeString(encodedCredentials)
	if err != nil {
		return "", false // 解码失败
	}

	// 解码后的格式是 "username:password"
	parts := strings.SplitN(string(decoded), ":", 2)
	if len(parts) != 2 {
		return "", false
	}

	username := parts[0]
	password := parts[1]

	// 检查用户名和密码是否有效
	if expectedPassword, ok := p.Credentials[username]; ok && expectedPassword == password {
		return username, true // 验证成功
	}

	return username, false // 验证失败
}

// handleConnection 处理单个客户端连接。
func (p *Proxy) handleConnection(clientConn net.Conn) {
	defer p.wg.Done()
	defer clientConn.Close()

	reader := bufio.NewReader(clientConn)
	req, err := http.ReadRequest(reader)
	if err != nil {
		if err != io.EOF {
			log.Printf("从 %s 读取 HTTP 请求失败: %v", clientConn.RemoteAddr(), err)
		}
		return
	}

	// 如果启用了身份验证，则进行检查
	if p.AuthEnabled {
		user, ok := p.checkAuth(req)
		if !ok {
			// 验证失败，发送 407 响应
			log.Printf("来自 %s 的身份验证失败 (用户: %s)", clientConn.RemoteAddr(), user)
			_, _ = fmt.Fprint(clientConn, "HTTP/1.1 407 Proxy Authentication Required\r\nProxy-Authenticate: Basic realm=\"Restricted\"\r\n\r\n")
			return
		}
		log.Printf("用户 [%s] (来自 %s) 身份验证成功", user, clientConn.RemoteAddr())
	}

	var destHost string
	if req.Method == http.MethodConnect {
		destHost = req.URL.Host
	} else {
		destHost = req.Host
		if !strings.Contains(destHost, ":") {
			destHost += ":80"
		}
	}

	if destHost == "" {
		log.Printf("无法从来自 %s 的请求中确定目标主机", clientConn.RemoteAddr())
		return
	}

	destConn, err := net.DialTimeout("tcp", destHost, 10*time.Second)
	if err != nil {
		log.Printf("连接目标服务器 %s 失败: %v", destHost, err)
		_, _ = fmt.Fprint(clientConn, "HTTP/1.1 502 Bad Gateway\r\n\r\n")
		return
	}
	defer destConn.Close()

	if req.Method == http.MethodConnect {
		_, err = fmt.Fprint(clientConn, "HTTP/1.1 200 Connection established\r\n\r\n")
		if err != nil {
			log.Printf("向客户端 %s 发送 200 OK 失败: %v", clientConn.RemoteAddr(), err)
			return
		}
	} else {
		err = req.Write(destConn)
		if err != nil {
			log.Printf("向目标服务器 %s 写入请求失败: %v", destHost, err)
			return
		}
	}

	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		_, _ = io.Copy(destConn, clientConn)
	}()

	go func() {
		defer wg.Done()
		_, _ = io.Copy(clientConn, destConn)
	}()

	wg.Wait()
}

func main() {
	// --- 配置区域 ---
	// 将此项设置为 true 来启用身份验证，设置为 false 则禁用
	const enableAuth = true

	// 如果启用身份验证，请在此处定义有效的用户名和密码
	validUsers := map[string]string{
		"user1": "pass1",
		"user2": "pass2",
		"admin": "aSecretPassword",
	}
	// --- 配置结束 ---

	proxy, err := NewProxy(":9191", enableAuth, validUsers)
	if err != nil {
		log.Fatalf("创建代理失败: %v", err)
	}

	if err := proxy.Start(); err != nil {
		log.Fatalf("启动代理失败: %v", err)
	}
}